/**
 * (c) Zoltan Farkas 2012
 */
package com.zoltran.stackmonitor;

import com.google.common.base.Function;
import java.io.FileWriter;
import java.io.IOException;
import java.io.Writer;
import java.lang.management.ManagementFactory;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;
import javax.annotation.concurrent.GuardedBy;
import javax.annotation.concurrent.ThreadSafe;
import javax.management.InstanceAlreadyExistsException;
import javax.management.MBeanRegistrationException;
import javax.management.MalformedObjectNameException;
import javax.management.NotCompliantMBeanException;
import javax.management.ObjectName;

/**
 * Utility that allow you to sample what the application is doing.
 * It generates a "Flame Graph" that allows you to quickly see you "heavy" 
 * operations.
 * 
 * You can use JConsole to control the sampling while your application is running.
 * 
 * By using a sampling approach you can choose your overhead.
 * (sampling takes about 0.5 ms, so the default of 10Hz will give you 0.5% overhead)
 * 
 * Collection is separated into CPU, WAIT and IO categories.
 * I felt that most important is to see what hogs your CPU because that is where,
 * you can most likely can do something about it.
 * 
 * @author zoly
 */
@ThreadSafe
@edu.umd.cs.findbugs.annotations.SuppressWarnings("I18N")
public class Sampler implements SamplerMBean {

    private volatile boolean stopped;
    private volatile long sampleTimeMillis;
    
    private final Object sampleSync = new Object();
    @GuardedBy("sampleSync")
    private SampleNode cpuSamples;
    @GuardedBy("sampleSync")
    private SampleNode waitSamples;
    @GuardedBy("sampleSync")
    private SampleNode ioSamples;
    @GuardedBy("this")
    private Thread samplingThread;
    
    private static final Set<Method> WAIT_METHODS = new HashSet();

    static {
        WAIT_METHODS.add(new Method(sun.misc.Unsafe.class, "park"));
        WAIT_METHODS.add(new Method(java.lang.Object.class, "wait"));
        WAIT_METHODS.add(new Method(java.lang.Thread.class, "sleep"));
    }
    private static final Set<Method> IO_METHODS = new HashSet();

    static {
        IO_METHODS.add(new Method(java.io.FileOutputStream.class, "writeBytes"));
        IO_METHODS.add(new Method(java.io.FileOutputStream.class, "open"));
        IO_METHODS.add(new Method(java.io.FileOutputStream.class, "openAppend"));
        IO_METHODS.add(new Method(java.io.FileDescriptor.class, "sync"));
        IO_METHODS.add(new Method(java.nio.MappedByteBuffer.class, "load0"));
        IO_METHODS.add(new Method(java.nio.MappedByteBuffer.class, "force0"));
        IO_METHODS.add(new Method("java.net.PlainSocketImpl", "socketAccept"));
        IO_METHODS.add(new Method("java.net.PlainSocketImpl", "socketConnect"));
    }

    private static boolean isWaitMethod(Method m) {
        return WAIT_METHODS.contains(m);
    }

    private static boolean isIOMethod(Method m) {
        return IO_METHODS.contains(m);
    }

    public Sampler() {
        this(100);
    }

    public Sampler(long sampleTimeMillis) {
        stopped = true;
        this.sampleTimeMillis = sampleTimeMillis;
        cpuSamples = null;
        waitSamples = null;

    }

    @PostConstruct
    public void init() throws MalformedObjectNameException, InstanceAlreadyExistsException, MBeanRegistrationException, NotCompliantMBeanException {
        ManagementFactory.getPlatformMBeanServer().registerMBean(this, new ObjectName("StackMonitor:name=StackSampler"));
    }

    @Override
    public synchronized void start() {
        if (stopped) {
            stopped = false;
            samplingThread = new Thread(new Runnable() {

                @SuppressWarnings("SleepWhileInLoop")
                public void run() {
                    while (!stopped) {
                        Map<Thread, StackTraceElement[]> stackDump = Thread.getAllStackTraces();
                        recordStackDump(stackDump);
                        try {
                            Thread.sleep(sampleTimeMillis);
                        } catch (InterruptedException ex) {
                            stopped = true;
                        }
                    }
                }
            }, "Stack Sampling Thread");
            samplingThread.start();
        } else {
            throw new IllegalStateException("Sampling can only be started once");
        }

    }

    private void recordStackDump(Map<Thread, StackTraceElement[]> stackDump) {
        synchronized (sampleSync) {
            for (Map.Entry<Thread, StackTraceElement[]> entry : stackDump.entrySet()) {
                StackTraceElement[] stackTrace = entry.getValue();
                if (stackTrace.length > 0 && !entry.getKey().equals(Thread.currentThread())) {
                    Method m = new Method(stackTrace[0]);
                    if (isWaitMethod(m)) {
                        if (waitSamples == null) {
                            waitSamples = new SampleNode(stackTrace, stackTrace.length - 1);
                        } else {
                            waitSamples.addSample(stackTrace, stackTrace.length - 1);
                        }

                    } else if (isIOMethod(m)) {
                        if (ioSamples == null) {
                            ioSamples = new SampleNode(stackTrace, stackTrace.length - 1);
                        } else {
                            ioSamples.addSample(stackTrace, stackTrace.length - 1);
                        }

                    } else {
                        if (cpuSamples == null) {
                            cpuSamples = new SampleNode(stackTrace, stackTrace.length - 1);
                        } else {
                            cpuSamples.addSample(stackTrace, stackTrace.length - 1);
                        }

                    }
                }
            }
        }
    }

    @Override
    public synchronized void generateHtmlMonitorReport(String fileName, int chartWidth) throws IOException {
        synchronized (sampleSync) {
            Writer writer = new FileWriter(fileName);
            try {
                writer.append("<html>");
                if (cpuSamples != null) {
                    writer.append("<h1>CPU stats</h1>");
                    generateCpuHtmlTable(writer, chartWidth);
                }
                if (ioSamples != null) {
                    writer.append("<h1>IO stats</h1>");
                    generateIOHtmlTable(writer, chartWidth);
                }
                if (waitSamples != null) {
                    writer.append("<h1>WAIT stats</h1>");
                    generateWaitHtmlTable(writer, chartWidth);
                }
                writer.append("</html>");

            } finally {
                writer.close();
            }
        }

    }

    public boolean generateCpuHtmlTable(Writer writer, int width) throws IOException {
        synchronized (sampleSync) {
            if (cpuSamples != null) {
                generateHtmlTable(writer, Method.ROOT, cpuSamples, width);
                return true;
            } else {
                return false;
            }
        }
    }

    public boolean generateWaitHtmlTable(Writer writer, int width) throws IOException {
        synchronized (sampleSync) {
            if (waitSamples != null) {
                generateHtmlTable(writer, Method.ROOT, waitSamples, width);
                return true;
            } else {
                return false;
            }
        }
    }

    public boolean generateIOHtmlTable(Writer writer, int width) throws IOException {
        synchronized (sampleSync) {
            if (ioSamples != null) {
                generateHtmlTable(writer, Method.ROOT, ioSamples, width);
                return true;
            } else {
                return false;
            }
        }
    }
    private static final String[] COLORS = {"#CCE01B",
        "#DDE01B", "#EEE01B", "#FFE01B", "#FFD01B",
        "#FFC01B", "#FFA01B", "#FF901B", "#FF801B",
        "#FF701B", "#FF601B", "#FF501B", "#FF401B"};

    private void generateHtmlTable(Writer writer, Method m, SampleNode node, int tableWidth) throws IOException {
        writer.append("<table border=\"0\" cellpadding=\"0\" cellspacing=\"0\" style=\"overflow:hidden;table-layout:fixed;width:").
                append(Integer.toString(tableWidth)).append("\"><tr><td title=\"");
        m.toWriter(writer);
        writer.append("\" style=\"vertical-align:top ;background:").
                append(COLORS[(int) (Math.random() * COLORS.length)]).append("\">");
        m.toWriter(writer);
        writer.append("</td></tr>\n");
        int totalSamples = node.getCount();

        Map<Method, SampleNode> subNodes = node.getSubNodes();
        if (subNodes != null) {
            writer.append("<tr><td><table border=\"0\" cellpadding=\"0\" cellspacing=\"0\" style=\"overflow:hidden;table-layout:fixed\"><tr>\n");
            for (Map.Entry<Method, SampleNode> entry : subNodes.entrySet()) {
                int width = entry.getValue().getCount() * tableWidth / totalSamples;
                writer.append("<td style=\"vertical-align:top; width:").append(Integer.toString(width)).append("px\">");
                generateHtmlTable(writer, entry.getKey(), entry.getValue(), width);
                writer.append("</td>");
            }
            writer.append("</tr></table></td></tr>");
        }
        writer.append("</table>\n");
    }

    @PreDestroy
    @Override
    public synchronized void stop() throws InterruptedException {
        stopped = true;
        samplingThread.join();
    }

    public void clear() {
        synchronized (sampleSync) {
            cpuSamples = null;
            waitSamples = null;
            ioSamples = null;
        }
    }

    public long getSampleTimeMillis() {
        return sampleTimeMillis;
    }

    public void setSampleTimeMillis(long sampleTimeMillis) {
        this.sampleTimeMillis = sampleTimeMillis;
    }

    public boolean isStopped() {
        return stopped;
    }

    public SampleNode applyOnCpuSamples(Function<SampleNode, SampleNode> predicate) {
        synchronized (sampleSync) {
            return cpuSamples = predicate.apply(cpuSamples);
        }
    }

    public SampleNode applyOnWaitSamples(Function<SampleNode, SampleNode> predicate) {
        synchronized (sampleSync) {
            return waitSamples = predicate.apply(waitSamples);
        }
    }

    public SampleNode applyOnIoSamples(Function<SampleNode, SampleNode> predicate) {
        synchronized (sampleSync) {
            return ioSamples = predicate.apply(ioSamples);
        }
    }
}
